package com.etl.dataflow.common.collection.queue;

import java.util.*;

/**
 * @author dx
 * @since 2022/6/21
 */
public class MinHeap<E extends Comparable<E>> {
    /**
     * 使用ArrayList的目的是无需关注动态扩缩容逻辑
     */
    private final ArrayList<E> data;

    public MinHeap(int capacity) {
        this.data = new ArrayList<>(capacity);
    }

    public MinHeap() {
        this.data = new ArrayList<>();
    }

    public MinHeap(List<E> list) {
        this();
        this.data.addAll(list);
        heapify();
    }

    public MinHeap(E[] arr) {
        this.data = asArrayList(arr);
        heapify();
    }

    /**
     * 将数组转换为ArrayList
     */
    private ArrayList<E> asArrayList(E[] arr) {
        ArrayList<E> ret = new ArrayList<>();
        Collections.addAll(ret, arr);
        return ret;
    }

    /**
     * 返回对中的元素个数
     */
    public synchronized int size() {
        return data.size();
    }

    /**
     * 判断堆是否为空
     */
    public synchronized boolean isEmpty() {
        return data.isEmpty();
    }

    /**
     * 根据传入的index，计算其父节点所在的下标
     */
    private int parent(int index) {
        if (index == 0) {
            throw new IllegalArgumentException("index-1 doesn't have parent.");

        }
        return (index - 1) / 2;
    }

    /**
     * 根据传入的index，计算其左子节点所在的下标
     */
    private int leftChild(int index) {
        return index * 2 + 1;
    }

    /**
     * 根据传入的index，计算其右子节点所在的下标
     */
    private int rightChild(int index) {
        return index * 2 + 2;
    }

    /**
     * 向堆中添加元素 e
     */
    public synchronized void add(E e, Comparator<E> keyComparator) {
        this.remove(e, keyComparator);
        data.add(e);
        if (data.size() == 1) {
            return;
        }
        heapify();
    }

    /**
     * 堆化
     */
    private void heapify() {
        // 最后一个非叶子节点的下标
        int lastNode = parent(data.size() - 1);
        for (int i = lastNode; i >= 0; i--) {
            // 从后往前依次堆化
            siftDown(i);
        }
    }

    /**
     * 获取堆顶元素
     */
    public synchronized E findMin() {
        if (isEmpty()) {
            return null;
        }
        return data.get(0);
    }

    /**
     * 从堆中取出元素，也就是取出堆顶元素
     */
    public synchronized E extractMin() {
        E ret = findMin();
        // 交换根节点与最后一个节点的位置
        Collections.swap(data, 0, data.size() - 1);
        // 删除最后一个节点
        data.remove(data.size() - 1);
        siftDown(0);
        return ret;
    }

    /**
     * 从上往下调整元素的位置，直到元素到达叶子节点或小于左右子节点
     */
    private void siftDown(int k) {
        // 左子节点大于size时就证明到底了
        while (leftChild(k) < data.size()) {
            int leftChildIndex = leftChild(k);
            int rightChildIndex = leftChildIndex + 1;
            int minChildIndex = leftChildIndex;

            // 左右子节点中最小的节点下标
            if (rightChildIndex < data.size() &&
                    isLessThan(rightChildIndex, leftChildIndex)) {
                minChildIndex = rightChildIndex;
            }

            // 小于最小的子节点证明 k 已经小于左右子节点，无需再继续下沉了
            if (data.get(k).compareTo(data.get(minChildIndex)) <= 0) {
                break;
            }

            // 否则，交换 k 与其最小子节点的位置，继续下沉
            Collections.swap(data, k, minChildIndex);
            k = minChildIndex;
        }
    }

    /**
     * 判断右子节点是否小于左子节点
     */
    private boolean isLessThan(int rightChildIndex, int leftChildIndex) {
        return data.get(rightChildIndex).compareTo(data.get(leftChildIndex)) < 0;
    }

    /**
     * 取出堆中的最大元素，并且替换成元素e
     */
    public synchronized E replace(E e) {
        E ret = findMin();
        // 替换堆顶元素
        data.set(0, e);
        siftDown(0);
        return ret;
    }

    public synchronized void remove(E e, Comparator<E> comparator) {
        Iterator<E> iterator = this.data.iterator();
        while (iterator.hasNext()) {
            E next = iterator.next();
            int compare = comparator.compare(next, e);
            if (compare == 0) {
                iterator.remove();
            }
        }
    }

    @Override
    public String toString() {
        return Arrays.toString(this.data.toArray(new Object[0]));
    }
}
